import torch
import numpy as np
import torch.nn as nn
import copy


def softmin(x, dim=-1, t=1):
    m = nn.Softmin(dim)

    return  m(x / t)
def compute_dice(input, target):
    eps = 0.0001
    # input 是经过了sigmoid 之后的输出。
    input = (input > 0.5).float()
    target = (target > 0.5).float()
    if target.sum() == 0:
        return 1.0

    inter = torch.sum(target.view(-1) * input.view(-1)) + eps

    # print(self.inter)
    union = torch.sum(input) + torch.sum(target) + eps

    t = (2 * inter.float()) / union.float()
    if int(t.item()) == 2:
        return 1.0
    return t.item()


def compute_dice_muti_class(input, target):

    ## 此时是多个类别应该得分开算dice
    class_num = input.shape[1]
    input = input.argmax(dim=1)
    # print(input.shape)
    dice = []
    for i in range(class_num):
        if i == 0:
            continue
        label = torch.zeros_like(target)
        input_single = torch.zeros_like(target)
        label[target == i] = 1
        input_single[input == i] = 1
        dice.append(compute_dice(input_single, label))
    return dice


def compute_iou(pred, label):
    # 计算iou loss
    inter = (pred * label).sum()
    union = torch.logical_or(pred, label)
    union = union.sum()

    return ((inter + 1e-6) /(union + 1e-6)).item()


def compute_3d_metric(pred_3d, label):
    ## 计算多个指标 3d
    # pred_3d_sig = (pred_3d_sig > 0.5).float()
    seg_inv, gt_inv = torch.logical_not(pred_3d), torch.logical_not(label)
    true_pos = float(torch.logical_and(pred_3d, label).sum())  # float for division
    true_neg = torch.logical_and(seg_inv, gt_inv).sum()
    false_pos = torch.logical_and(pred_3d, gt_inv).sum()
    false_neg = torch.logical_and(seg_inv, label).sum()

    # 然后根据公式分别计算出这几种指标
    prec = true_pos / (true_pos + false_pos + 1e-6)
    rec = true_pos / (true_pos + false_neg + 1e-6)
    specificity = true_neg / (true_neg + false_neg + 1e-6)

    return prec.item(), rec.item(), specificity.item()


def segmenation_metric(pred, target):
    class_num = pred.shape[1]
    pred = pred.argmax(dim=1)
    # print(input.shape)
    metric = []

    for i in range(class_num):
        if i == 0:
            continue
        label = torch.zeros_like(target)
        input_single = torch.zeros_like(target)
        label[target == i] = 1
        input_single[pred == i] = 1
        dice = compute_dice(input_single, label)
        prec, rec, specificity = compute_3d_metric(input_single, label)
        iou = compute_iou(input_single, label)

        metric.append([dice, prec, rec, specificity, iou])
    # print(metric)
    return metric


class SoftmaxLoss(nn.Module):
    def __init__(self, reduction="mean"):
        super(SoftmaxLoss, self).__init__()
        self.reduction = reduction
    def forward(self, output, target):
        return softmax_loss(output, target, reduction=self.reduction)


class SoftmaxLossHuanhu(nn.Module):
    def __init__(self, reduction="mean"):
        super(SoftmaxLossHuanhu, self).__init__()
        self.reduction = reduction

    def forward(self, output, target):
        return softmax_loss_huanhu(output, target, reduction=self.reduction)


def softmax_loss(output, target, reduction="mean"):
    device = output.device
    pos_1 = torch.eq(target, 1).float()
    pos_2 = torch.eq(target, 2).float()
    pos_3 = torch.eq(target, 3).float()
    neg = torch.eq(target, 0).float()
    all_num = pos_1.sum() + pos_2.sum() + pos_3.sum() + neg.sum()

    pos_1_ratio = 1 - (pos_1.sum() / all_num)
    pos_2_ratio = 1 - (pos_2.sum() / all_num)
    pos_3_ratio = 1 - (pos_3.sum() / all_num)
    neg_ratio = 1 - (neg.sum() / all_num)
    weights = torch.tensor([neg_ratio, pos_1_ratio, pos_2_ratio, pos_3_ratio], device=device)

    # weights = pos_1_ratio * pos_1 + pos_2_ratio * pos_2 + pos_3_ratio * pos_3 + neg_ratio * neg
    return torch.nn.functional.cross_entropy(output, target, weights, reduction=reduction)

def softmax_loss_huanhu(output, target, reduction="mean"):
    device = output.device
    pos_1 = torch.eq(target, 1).float()

    neg = torch.eq(target, 0).float()
    all_num = pos_1.sum() + neg.sum()

    pos_1_ratio = 1 - (pos_1.sum() / all_num)
    neg_ratio = 1 - (neg.sum() / all_num)
    weights = torch.tensor([neg_ratio, pos_1_ratio], device=device)
    print(weights)
    # weights = pos_1_ratio * pos_1 + pos_2_ratio * pos_2 + pos_3_ratio * pos_3 + neg_ratio * neg
    return torch.nn.functional.cross_entropy(output, target, weights, reduction=reduction)

def infer2d_uncer(image3d, network):
    # image3d: batch, channels, D, W, H
    index = 0
    res_uncer_map = None
    res_learn_ratio_2d = None
    res_pred = None
    for y in torch.split(image3d, 1, dim=2):
        if index == 0:
            y = y.squeeze(2)
            pred_2d, pred_2d_uncer, learn_ratio = network.uncer_supervise_v2(y)
            res_pred = pred_2d.unsqueeze(2)
            res_uncer_map = pred_2d_uncer.unsqueeze(2)
            # res_learn_ratio_2d = learn_ratio.unsqueeze(dim=1)
        else:
            y = y.squeeze(2)
            pred_2d, pred_2d_uncer, learn_ratio = network.uncer_supervise_v2(y)
            res_pred = torch.cat((res_pred, pred_2d.unsqueeze(2)), dim=2)
            res_uncer_map = torch.cat((res_uncer_map, pred_2d_uncer.unsqueeze(2)), dim=2)
            # res_learn_ratio_2d = torch.cat((res_learn_ratio_2d, learn_ratio.unsqueeze(dim=1)), dim=1)

        index += 1
    return res_pred, res_uncer_map, None

def infer_25_uncer(image3d, net_25d):
    device = image3d.device
    batch, channel, d, h, w = image3d.shape
    index = 0
    res_pred = None
    res_uncer_map = None
    res_learn_ratio_25d = None

    for y in torch.split(image3d, 1, dim=2):
        y = y.squeeze(dim=2) # 去掉deep 维度 在channel维度进行cat
        if index == 0:
            y_in = torch.zeros((batch, 3, h, w), device=device)
            y = torch.cat([y, image3d[:, :, index + 1]], dim=1)
            y_in[:, 1:] = copy.deepcopy(y)
            del y
            pred_25d, pred_25d_uncer, learn_ratio_25d = net_25d.uncer_supervise_v2(y_in)
            res_uncer_map = pred_25d_uncer.unsqueeze(dim=2)
            res_pred = pred_25d.unsqueeze(2)
            # res_learn_ratio_25d = learn_ratio_25d.unsqueeze(dim=1)

        elif index == d - 1:
            y_in = torch.zeros((batch, 3, h, w), device=device)
            y = torch.cat([y, image3d[:, :, index - 1]], dim=1)
            y_in[:, :2] = copy.deepcopy(y)
            del y
            pred_25d, pred_25d_uncer, learn_ratio_25d = net_25d.uncer_supervise_v2(y_in)
            res_pred = torch.cat([res_pred, pred_25d.unsqueeze(dim=2)], dim=2)
            res_uncer_map = torch.cat([res_uncer_map, pred_25d_uncer.unsqueeze(dim=2)], dim=2)

            # res_learn_ratio_25d = torch.cat([res_learn_ratio_25d, learn_ratio_25d.unsqueeze(dim=1)], dim=1)

        else:
            y = torch.cat([image3d[:, :, index - 1], y, image3d[:, :, index + 1]], dim=1)
            pred_25d, pred_25d_uncer, learn_ratio_25d = net_25d.uncer_supervise_v2(y)
            res_pred = torch.cat((res_pred, pred_25d.unsqueeze(2)), dim=2)
            res_uncer_map = torch.cat((res_uncer_map, pred_25d_uncer.unsqueeze(2)), dim=2)
            # res_learn_ratio_25d = torch.cat([res_learn_ratio_25d, learn_ratio_25d.unsqueeze(dim=1)], dim=1)

        index += 1
    return res_pred, res_uncer_map, None


def sigmoid_rampup(current, rampup_length):
    """Exponential rampup from https://arxiv.org/abs/1610.02242"""
    if rampup_length == 0:
        return 1.0
    else:
        current = np.clip(current, 0.0, rampup_length)
        phase = 1.0 - current / rampup_length
        return float(np.exp(-5.0 * phase * phase))



# def infer2d(image3d, network):
#     # image3d: batch, channels, D, W, H
#     index = 0
#     res_img = None
#     for y in torch.split(image3d, 1, dim=2):
#         if index == 0:
#             y = y.squeeze(2)
#             pred_2d = network(y)
#             res_img = pred_2d.unsqueeze(2)
#         else:
#             y = y.squeeze(2)
#             pred_2d = network(y)
#             res_img = torch.cat((res_img, pred_2d.unsqueeze(2)), dim=2)
#         index += 1
#     return res_img
#
# def infer_25(image3d, net_25d, device="cpu"):
#     device = image3d.device
#     batch, channel, d, h, w = image3d.shape
#     index = 0
#     res_img = None
#     for y in torch.split(image3d, 1, dim=2):
#         y = y.squeeze(dim=2) # 去掉deep 维度 在channel维度进行cat
#         if index == 0:
#             y_in = torch.zeros((batch, 3, h, w), device=device)
#             y = torch.cat([y, image3d[:, :, index + 1]], dim=1)
#             y_in[:, 1:] = copy.deepcopy(y)
#             del y
#             pred_25d = net_25d(y_in)# (batch, 1, h, w)
#             res_img = pred_25d.unsqueeze(dim=2) ## (batch, 1, 1, h, w)
#         elif index == d - 1:
#             y_in = torch.zeros((batch, 3, h, w), device=device)
#             y = torch.cat([y, image3d[:, :, index - 1]], dim=1)
#             y_in[:, :2] = copy.deepcopy(y)
#             del y
#             pred_25d = net_25d(y_in)
#             res_img = torch.cat([res_img, pred_25d.unsqueeze(dim=2)], dim=2)
#         else:
#             y = torch.cat([image3d[:, :, index - 1], y, image3d[:, :, index + 1]], dim=1)
#             pred_25d = net_25d(y) ##( batch, 3, h, w) -> (batch, 4, h, w)
#             ## (batch, 4, d, h, w)
#             res_img = torch.cat((res_img, pred_25d.unsqueeze(2)), dim=2)
#         index += 1
#     return res_img
